#!/usr/bin/env python3

import readline
from dataclasses import dataclass
from typing import Callable, List, Dict, Optional
import time
import traceback

# split a command into tokens.
# It's similar to shlex.split, but handles ';' and '|'.
def split_command(s: str) -> "list[str]":
    result = []
    # states
    NORMAL = 1
    IN_QUOTE = 2
    t = ""
    q = ""
    state = NORMAL
    
    for c in s:
        if state == NORMAL:
            if c in [' ', '\t']:
                if t:
                    result.append(t)
                    t = ""
            else:
                if c == '|' or c == ';':
                    if t:
                        result.append(t)
                    result.append(c)
                    t = ""
                elif c in ['"', "'"]:
                    state = IN_QUOTE
                    q = c
                else:
                    t += c
        elif state == IN_QUOTE:
            if c == q:
                state = NORMAL
                result.append(t)
                t = ""
                q = ""
            else:
                t += c
                
    if t:
        result.append(t)
    
    return result


@dataclass
class Command:
    """
    A command with a name, help text, and an action to perform.
    """
    command: str
    help: str
    action: Callable[[list[str]], None] = None
    complete: Callable[[List[str]], List[str]] = None

    def next_arg(self, args: List[str]) -> str:
        """
        Compare the input arguments with the command's expected arguments.
        Returns None if no match is found or if there are too many arguments.
        Otherwise, returns the next expected argument.
        Args:
            args: The input arguments to compare with the command's expected arguments.
        Returns:
            The next expected argument or None if no match is found.
        """
        items = split_command(self.command)
        for i, arg in enumerate(args):
            item = items[i]
            if item.startswith('['):
                # this can match anything
                continue
            if not item.startswith(arg):
                return None
        if len(args) < len(items):
            return items[len(args)]
        else:
            return None
        
    def matches_args(self, args: List[str]) -> bool:
        """
        Check if the input arguments match the command's expected arguments.
        Args:
            args: The input arguments to check.
        Returns:
            True if the input arguments match the command's expected arguments, False otherwise.
        """
        items = split_command(self.command)
        if len(items) != len(args):
            return False
        for i, arg in enumerate(args):
            item = items[i]
            if item.startswith('['):
                # this can match anything
                continue
            if not item.startswith(arg):
                return False
        return True


class CommandManager:
    """ A manager for commands. """
    def __init__(self):
        self.commands = [
            Command("modules", "List all the modules", self.action_modules),
            Command("config [module]", "Config a module.", self.action_config, self.complete_config),
            Command("set [variable] [value]", "Set a variable to a value.", self.action_set, self.complete_set),
            Command("list", "List all items in a module.", self.action_list),
            Command("exit", "Exits current scope", self.action_exit),
            Command("help", "Displays help information.", self.action_help),
        ]
        self.modules = {
            "system": {
                "hostname": "host1",
                "os": "Linux",
                "kernel": "5.4.0-72-generic"
            },
            "network": {
                "ip": "192.168.0.100",
                "gateway": "192.168.0.1",
                "dns": "8.8.8.8"
            }
        }
        # currently selected module
        self.module = None

    def action_modules(self, args: list[str]) -> None:
        for module in self.modules.keys():
            print(module)
    
    def action_config(self, args: list[str]) -> None:
        if len(args) < 2:
            print("Usage: config [module]")
            return
        module = args[1]
        if module not in self.modules:
            print(f"Module '{module}' not found.")
            return
        print(f"Configuring module '{module}'")
        self.module = module

    def complete_config(self, args: List[str]) -> List[str]:
        if len(args) == 1:
            return [module for module in self.modules.keys()]
        return []
    
    def action_set(self, args: list[str]) -> None:
        if len(args) < 3:
            print("Usage: set [key] [value]")
            return
        if not self.module:
            print("No module selected. Use 'config' to select a module.")
            return
        key = args[1]
        value = " ".join(args[2:])
        print(f"Setting {key} to {value}")
        self.modules[self.module][key] = value

    def complete_set(self, args: List[str]) -> List[str]:
        if not self.module:
            return []
        if len(args) == 1:
            return [key for key in self.modules[self.module].keys()]
        if len(args) == 2:
            key = args[1]
            return [self.modules[self.module][key]]
        return []
    
    def action_list(self, args: list[str]) -> None:
        # list all the keys for the current module
        print(f"Values in module '{self.module}':\n")
        for key, value in self.modules[self.module].items():
            print(f"{key.ljust(20)}: {value}")

    def action_exit(self, args: list[str]) -> None:
        if self.module:
            self.module = None
            return
        raise SystemExit()

    def action_help(self, args: list[str]) -> None:
        print("Available commands:\n")
        max_len = max(len(command.command) for command in self.commands) + 2
        for command in self.commands:
            print(f"{command.command.ljust(max_len)}: {command.help}")

    def get_prompt(self) -> str:
        if not self.module:
            return "> "
        return f"config {self.module}# "
    
    def get_candidate_commands(self, args: List[str]) -> List[str]:
        pass
    
    def complete(self, text: str, state: int) -> Optional[str]:
        """
        Complete the current word.
        """
        # get current line
        line = readline.get_line_buffer()
        # split the line
        words = split_command(line)
        if words and not line.endswith(' '):
            # prefix is the last word
            prefix = words[-1]
            words = words[:-1]
        else:
            prefix = ''
        candidates = []
        for command in self.commands:
            # check whether the command can match the word list
            candidate = command.next_arg(words)
            if candidate:
                if candidate.startswith('['):
                    possibilities = command.complete(words)
                    if possibilities:
                        candidates += [x for x in possibilities if x.startswith(prefix)]
                elif candidate.startswith(prefix):
                    candidates.append(candidate)
        candidates.sort()
        if state < len(candidates):
            return candidates[state] + " "
        return None
    
    def execute(self, line: str) -> None:
        args = split_command(line)
        if not args:
            return
        for command in self.commands:
            if command.matches_args(args):
                command.action(args)
                return
        print(f"Unknown command: {args}")


class CommandLine:
    def __init__(self, ctx: CommandManager):
        self.ctx = ctx

    def loop(self):
        try:
            readline.parse_and_bind("tab: complete")
            readline.parse_and_bind("?: possible-completions")
            delims = readline.get_completer_delims()
            # dont use '-' for delim
            delims = delims.replace('-', '')
            readline.set_completer_delims(delims)
            readline.set_completer(self.ctx.complete)
        except:
            pass

        line = ""
        ctrlCTime = None
        while True:
            try:
                prompt = self.ctx.get_prompt()
                try:
                    line = input(prompt)
                    self.ctx.execute(line)
                except KeyboardInterrupt:
                    # Press Ctrl+C twice in 2 seconds to exit
                    now = time.time()
                    if ctrlCTime is None or now - ctrlCTime > 2.0:
                        ctrlCTime = now
                        print("Press Ctrl+C again to exit.")
                    else:
                        raise SystemExit(1)
            except (EOFError, SystemError, SystemExit, KeyboardInterrupt):
                print("\nBye")
                break
            except:
                print(f"Failed to run {line}")
                traceback.print_exc()

def main():
    ctx = CommandManager()
    cli = CommandLine(ctx)
    cli.loop()


if __name__ == "__main__":
    main()
